#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
计算所有患者的平均指标并添加到experiment_summary.json文件中
"""
import json
import os

def calculate_overall_metrics(summary_file):
    """
    计算所有患者的平均指标并添加到实验摘要文件中
    
    Args:
        summary_file: 实验摘要文件路径
    """
    try:
        # 读取现有的实验摘要文件
        with open(summary_file, 'r', encoding='utf-8') as f:
            experiment_summary = json.load(f)
        
        # 获取结果摘要列表
        results_summary = experiment_summary.get('results_summary', [])
        total_patients = len(results_summary)
        
        if total_patients == 0:
            print("没有找到患者数据，无法计算平均指标")
            return
        
        print(f"发现{total_patients}个患者的实验结果，开始计算平均指标...")
        
        # 初始化指标累加器
        total_metrics = {
            'bertscore': 0.0,
            'rouge-1': 0.0,
            'rouge-2': 0.0,
            'rouge-l': 0.0,
            'rouge-1_p': 0.0,
            'rouge-1_r': 0.0,
            'rouge-2_p': 0.0,
            'rouge-2_r': 0.0,
            'rouge-l_p': 0.0,
            'rouge-l_r': 0.0,
            'meteor': 0.0,
            'length_ratio': 0.0
        }
        
        total_bleu = {
            'bleu-1': 0.0,
            'bleu-2': 0.0,
            'bleu-3': 0.0,
            'bleu-4': 0.0
        }
        
        # 遍历所有患者的指标，累加计算
        for patient_summary in results_summary:
            metrics = patient_summary.get('metrics_summary', {})
            
            # 累加非嵌套指标
            for key in total_metrics:
                if key in metrics:
                    total_metrics[key] += metrics[key]
            
            # 处理嵌套的BLEU指标
            if 'bleu' in metrics and isinstance(metrics['bleu'], dict):
                for key in total_bleu:
                    if key in metrics['bleu']:
                        total_bleu[key] += metrics['bleu'][key]
        
        # 计算平均值
        overall_metrics = {}
        for key in total_metrics:
            overall_metrics[key] = round(total_metrics[key] / total_patients, 4)
        
        # 计算BLEU指标的平均值
        overall_bleu = {}
        for key in total_bleu:
            overall_bleu[key] = round(total_bleu[key] / total_patients, 4)
        
        # 合并BLEU指标
        overall_metrics['bleu'] = overall_bleu
        
        # 添加所有患者的总平均指标到实验摘要
        experiment_summary['overall_metrics_average'] = overall_metrics
        
        # 递归处理结果中的所有浮点数，确保它们以小数形式显示
        def process_floats_for_summary(obj):
            if isinstance(obj, float):
                # 对于接近0的极小值，直接显示为0.0
                if abs(obj) < 1e-10:
                    return 0.0
                # 对于较大的数值或需要保留小数的数值，格式化为字符串
                # 这样可以避免JSON序列化时自动使用科学计数法
                # 先尝试格式化为10位小数，如果是整数则去掉小数点
                formatted = "{0:.10f}".format(obj)
                if formatted.endswith('.0000000000'):
                    return float(formatted[:-10] + '.0')
                else:
                    # 去掉末尾的0和可能的小数点
                    formatted = formatted.rstrip('0').rstrip('.') if '.' in formatted else formatted
                    return float(formatted)
            elif isinstance(obj, dict):
                return {k: process_floats_for_summary(v) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [process_floats_for_summary(item) for item in obj]
            else:
                return obj
        
        # 处理结果中的所有浮点数
        processed_summary = process_floats_for_summary(experiment_summary)
        
        # 写回更新后的实验摘要文件
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(processed_summary, f, ensure_ascii=False, indent=2, allow_nan=False)
        
        print(f"所有患者的平均指标已成功添加到文件: {summary_file}")
        print("平均指标计算结果:")
        for key, value in overall_metrics.items():
            if key == 'bleu':
                print(f"  {key}:")
                for bleu_key, bleu_value in value.items():
                    print(f"    {bleu_key}: {bleu_value}")
            else:
                print(f"  {key}: {value}")
        
    except Exception as e:
        print(f"计算平均指标时发生错误: {str(e)}")

if __name__ == "__main__":
    # 实验摘要文件路径
    summary_file = os.path.join(os.path.dirname(__file__), '../results/multi_turn/experiment_summary.json')
    
    # 调用函数计算平均指标
    calculate_overall_metrics(summary_file)